import torch

from PIL import Image
from matplotlib import pyplot as plt
from torchvision.transforms import Compose, ToTensor, Grayscale
import torch.nn.functional as func
import torch.nn as nn
import torch.optim as optim

import lithosim_cuda as ilt
import ilt_utils as ut
import epe_checker as ev

DOSE_NOM = 1.0
DOSE_MAX = 1.02
DOSE_MIN = 0.98

def forward(mask, threshold, kernel_focus, kernel_defcous, weight_focus, weight_defocus, wafer_output_path, save_bin_wafer_image):
    aerial_nom, binary_wafer_nom = ilt.lithosim(mask, threshold, kernel_focus, weight_focus, wafer_output_path, save_bin_wafer_image,
                                                None, None, DOSE_NOM, True)
    aerial_max, binary_wafer_max = ilt.lithosim(mask, threshold, kernel_focus, weight_focus, wafer_output_path, save_bin_wafer_image,
                                                None, None, DOSE_MAX, True)
    aerial_nom, binary_wafer_min = ilt.lithosim(mask, threshold, kernel_defcous, weight_defocus, wafer_output_path, save_bin_wafer_image,
                                                None, None, DOSE_MIN, True)
    l2loss = func.mse_loss(binary_wafer_nom, mask.squeeze(0), reduction="sum")
    pvband = torch.sum(binary_wafer_max.bool() != binary_wafer_min.bool())
    print(f'l2loss: {l2loss.item()}')
    print(f'pvband: {pvband.item()}')

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为 SimHei
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

    plt.imshow(aerial_nom.squeeze().cpu().numpy())
    plt.title('nominal空间像')
    plt.show()

    plt.imshow(binary_wafer_nom.squeeze().cpu().numpy())
    plt.title('nominal光刻胶像')
    plt.show()

    return binary_wafer_nom, binary_wafer_max, binary_wafer_min

if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 设置字体为支持中文的字体，例如 SimHei（黑体）
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为 SimHei
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

    # mask_path = 'ICCAD2013/png/target1.png'
    mask_path = '../ICCAD2013/png/target1.png'
    # mask_path = r'E:\Papers\0 ILT\代码\传统方法\ICCAD2013\png\target1.png'
    mask = ilt.load_image(mask_path).to(device)
    # avgpool_size = 4  # 降采样因子
    # avg_layer4 = nn.AvgPool2d(avgpool_size, stride=avgpool_size)
    # mask_lowres4 = avg_layer4(mask)  # [1,512,512]
    # target_lowres4 = avg_layer4(mask)  # [1,512,512]
    #
    # avgpool_size2 = 2  # 降采样因子
    # avg_layer2 = nn.AvgPool2d(avgpool_size2, stride=avgpool_size2)
    # mask_lowres2 = avg_layer2(mask)
    # target_lowres2 = avg_layer2(mask)

    threshold = 0.225
    save_bin_wafer_image = False
    kernel = torch.load('../kernel_neuralilt/kernel_focus_tensor.pt', map_location=torch.device("cuda"))
    # kernel_lowres2 = torch.load('kernel_neuralilt/kernel_focus_tensor_lowres2.pt', map_location=torch.device("cuda"))
    # kernel_lowres4 = torch.load('kernel_neuralilt/kernel_focus_tensor_lowres4.pt', map_location=torch.device("cuda"))
    weight = torch.load('../kernel_neuralilt/weight_focus_tensor.pt', map_location=torch.device("cuda"))
    dose_nom = 1.0
    dose_max = 1.02
    dose_min = 0.98

    wafer_output_path = 'tmp/ineur_wafer1.png'
    aerial_nom, binary_wafer_nom = ilt.lithosim(mask, threshold, kernel, weight, None, False,
                                               None, None,DOSE_NOM,True)

    # aerial_nom_lowres4, binary_wafer_nom_lowres4 = ilt.lithosim(mask_lowres4, threshold, kernel_lowres4, weight, None, False,
    #                                             None, None, DOSE_NOM, True)
    #
    # aerial_nom_lowres2, binary_wafer_nom_lowres2 = ilt.lithosim(mask_lowres2, threshold, kernel_lowres2, weight, None, False,
    #                                             None, None, DOSE_NOM, True)

    plt.imshow(aerial_nom.squeeze().cpu().numpy())
    plt.title('原图nominal空间像')
    plt.colorbar()
    plt.show()

    # plt.imshow(aerial_nom_lowres4.squeeze().cpu().numpy())
    # plt.title('4倍降采样nominal空间像')
    # plt.colorbar()
    # plt.show()
    #
    # plt.imshow(aerial_nom_lowres2.squeeze().cpu().numpy())
    # plt.title('2倍降采样nominal空间像')
    # plt.colorbar()
    # plt.show()

    plt.imshow(binary_wafer_nom.squeeze().cpu().numpy())
    plt.title('原图nominal光刻胶像')
    plt.colorbar()
    plt.show()


    # kernel_f = torch.load('kernel_neuralilt/kernel_focus_tensor.pt', map_location=torch.device("cuda"))
    # kernel_df = torch.load('kernel_neuralilt/kernel_defocus_tensor.pt', map_location=torch.device("cuda"))
    # weight_f = torch.load('kernel_neuralilt/weight_focus_tensor.pt', map_location=torch.device("cuda"))
    # weight_df = torch.load('kernel_neuralilt/weight_defocus_tensor.pt', map_location=torch.device("cuda"))
    #
    # bwafer_nom, bwafer_max, bwafer_min = forward(mask, threshold, kernel_f, kernel_df, weight_f, weight_df, wafer_output_path, save_bin_wafer_image)
    #
    # # fig, axes = plt.subplots(2, 2)
    # # ax1, ax2, ax3, ax4 = axes.flatten()
    # # ax1.imshow(mask.squeeze(0).cpu())
    # # ax1.set_title('Mask')
    # #
    # # ax2.imshow(bwafer_nom.cpu())
    # # ax2.set_title('Binary Wafer Nom')
    # #
    # # ax3.imshow(bwafer_max.cpu())
    # # ax3.set_title('Binary Wafer Max')
    #
    # ax4.imshow(bwafer_min.cpu())
    # ax4.set_title('Binary Wafer Min')
    #
    # plt.show()

    # resist_nom = ev.sigmoid_ilt_z(aerial_nom)
    #
    # ax2.imshow(resist_nom.squeeze(0).cpu())
    # ax2.set_title('Resist Nom')
    # plt.show()

    # sigmod_mask = ev.sigmoid_ilt_mask(mask)
    # diff_mask = mask.squeeze(0) - sigmod_mask.squeeze(0)
    # ax2.imshow(diff_mask.cpu())
    # ax2.set_title('Diff Mask')
    # plt.show()




